import importlib.abc
import importlib.util
import sys
import os
import logging
import tempfile
import atexit


class VllmPatchFinder(importlib.abc.MetaPathFinder):
    def __init__(self, patch_info):
        self.patch_info = patch_info
        self.patch_folder = tempfile.TemporaryDirectory()
        atexit.register(self.patch_folder.cleanup)
        
    def get_py_file(self, fullname, path):
        module_name = fullname.split('.')[-1]
        module_file_path = os.path.join(path, f"{module_name}.py")
        module_init_path = os.path.join(path, module_name, "__init__.py")
        if os.path.exists(module_file_path):
            return module_file_path
        elif os.path.exists(module_init_path):
            return module_init_path
        else:
            return None
    
    def get_patch_info(self, fullname):
        return self.patch_info.get(fullname)
        
    def patch(self, fullname, source_file, patch_info):
        from patch import apply, fromstring
        patch_set = fromstring(patch_info)

        with open(source_file, 'r') as f:
            original_content = f.read()

        try:
            patched_content = apply(patch_set, original_content)
            patched_file = os.path.join(self.patch_folder, f"{fullname}.py")
            with open(patched_file, 'w') as f:
                f.write(patched_content)
            print(f"Patch {patched_content} applied successfully!")
            return patched_file
        except Exception as e:
            print(f"Failed to apply patch {patched_content} : {e}")
            return None
    
    def find_spec(self, fullname, path, target=None):
        if not fullname.startswith('vllm.'):
            return None
        
        patch_info =  self.get_patch_info(fullname)
        if not patch_info:
            return None
        
        source_file = self.get_py_file(fullname, path)
        if not source_file:
            return None
        
        patched_file = self.patch(fullname, source_file, patch_info)

        return importlib.util.spec_from_file_location(fullname, patched_file)

def patch_vllm(vllm_version):
    patch_file_path = os.path.join(os.path.dirname(__file__), 'patchs', vllm_version)
    if not os.path.exists(patch_file_path):
        logging.error(f"Patch file {patch_file_path} not found. maybe version {vllm_version} not support")
        return
    
    with open(patch_file_path, 'r') as f:
        original_content = f.read()
    
    patch_info = {}
    for diff in original_content.split("diff --git "):
        file_path = diff.split(" ")[0]
        if not file_path.endswith(".py"):
            continue
        file_path = os.path.basename(file_path)
        file_path.replace("/", ".")
        if file_path.endswith(".__init__"):
            file_path = os.path.dirname(file_path)
            
        patch_info[file_path] = diff
    
    # 将自定义查找器插入 sys.meta_path
    finder = VllmPatchFinder(patch_info)
    sys.meta_path.insert(0, finder)